package com.li.config;

import com.li.common.ModelName;
import com.li.dao.RedisChatMemoryStore;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import jakarta.annotation.Resource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

@Component
public class RagAssistantFactory {

    private final ChatModel chatModel;
    private final StreamingChatModel streamingChatModel;
    private final EmbeddingStore<TextSegment> embeddingStore;
    @Resource
    private RedisTemplate<String, String> redisTemplate;

    public RagAssistantFactory(
            ChatModel chatModel,
            StreamingChatModel streamingChatModel,
            EmbeddingStore<TextSegment> embeddingStore) {
        this.chatModel = chatModel;
        this.streamingChatModel = streamingChatModel;
        this.embeddingStore = embeddingStore;
    }

    public AiConfig.Assistant createAssistant(Integer id) {
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .id(id)
                .maxMessages(10)
                .chatMemoryStore(new RedisChatMemoryStore(redisTemplate,"memoryId"
                        , 60 * 60 * 24 *5))
                .build();

        QwenEmbeddingModel embeddingModel = QwenEmbeddingModel.builder()
                .modelName(ModelName.TEXT_EMBEDDING_V3)
                .apiKey(System.getenv("TONGYI_KEY"))
                .build();

        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(5) //返回相似度最多的5个结果
                .minScore(0.6) //最小匹配得分
                .build();

        return AiServices.builder(AiConfig.Assistant.class)
                .chatModel(chatModel)
                .streamingChatModel(streamingChatModel)
                .contentRetriever(contentRetriever)
                .chatMemoryProvider(memoryId -> chatMemory)
                .build();
    }
}
